package com.atguigu.aiproject.aitools;

import com.atguigu.aiproject.service.AdminService;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import org.springframework.ai.document.Document;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Component
public class UserTools {

    private final AdminService adminService;
    private final VectorStore creationVectorStore;
    private final ObjectMapper objectMapper = new ObjectMapper();

    public UserTools(AdminService adminService, VectorStore creationVectorStore) {
        this.adminService = adminService;
        this.creationVectorStore = creationVectorStore;
        objectMapper.enable(SerializationFeature.INDENT_OUTPUT); // 美化输出
    }

    @Tool(description = "SQL查询工具，只能执行select语句，执行并返回查询结果")
    public String query(@ToolParam(description = "sql查询语句") String sql) {
        System.out.println("查询：" + sql);
        List<Map<String, Object>> result;
        try {
            result = adminService.selectSql(sql);
        } catch (Exception e) {
            e.printStackTrace();
            return sql + "执行失败\n" + e.getMessage();
        }
        // 将结果转换为JSON字符串
        try {
            System.out.println(objectMapper.writeValueAsString(result));
            return objectMapper.writeValueAsString(result);
        } catch (JsonProcessingException e) {
            // 如果JSON转换失败，回退到toString()
            e.printStackTrace();
            return result.toString();
        }
    }

    @Tool(description = "作品查询工具，返回与输入内容关联度较高的所有作品ID")
    public String vectorStoreQuery(@ToolParam(description = "查询内容") String content) {
        SearchRequest searchRequest = new SearchRequest.Builder()
                .query(content)
                .similarityThreshold(0.45f)
                .topK(150) // todo topK应设置为数据库所有作品数量
                .build();
        List<Document> documents = creationVectorStore.similaritySearch(searchRequest);
        List<Integer> creationIds = new ArrayList<>();
        for (Document document : documents) {
            String creationId = (String) document.getMetadata().get("作品ID");
            creationIds.add(Integer.parseInt(creationId));
        }
        // 将结果转换为JSON字符串
        try {
            System.out.println(objectMapper.writeValueAsString(creationIds));
            return objectMapper.writeValueAsString(creationIds);
        } catch (JsonProcessingException e) {
            // 如果JSON转换失败，回退到toString()
            e.printStackTrace();
            return creationIds.toString();
        }
    }
}
